# coding: utf-8
import nltk
import os
import codecs
import sys
import random
from scrape.mongo import MongoConn
import re
import json
import sys

if len(sys.argv) != 2:
    print('usage: python gen_xrs_json.py [primary|middle|high]')
else:
    target = sys.argv[1].strip()
    assert target in ['primary', 'middle', 'high']


condition = {'difficulty': target}
jfilename = 'data/{}/all.json'.format(target)
filename = jfilename.replace('.json', '.txt')
m = MongoConn()

if os.path.exists(os.path.dirname(filename)) is False:
    os.makedirs(os.path.dirname(filename))


def fill_in_blank(question_json):
    answers = question_json['choices'][question_json['answer']]
    question = question_json['question']
    for answer in answers:
        question = question.replace('<BLANK>', answer)
    return question


with codecs.open(filename, 'w', encoding='utf-8') as f:
    with codecs.open(jfilename, 'w', encoding='utf-8') as jf:
        results = m.find(condition)
        random.shuffle(results)
        for question_json in results:
            try:
                question_json['question'] = re.sub(r'(?<!\w)-', ' - ', question_json['question'])  # 针对一些 - ，做特殊处理
                assert len(question_json['choices']) != 1  # 这些是错误抓取
                line = fill_in_blank(question_json)
                line = ' '.join(nltk.tokenize.word_tokenize(line))
                jf.write(json.dumps(question_json))
                jf.write('\n')
                f.write(line)
                f.write('\n')
            except Exception, e:
                print(e)
                print(question_json)
